import gradio as gr
from sklearn.datasets import load_iris
from sklearn.svm import SVC
import pickle


def svc_demo(C, kernel, degree, save):
    X, y = load_iris(return_X_y=True)
    svc = SVC(C=C, kernel=kernel, degree=degree)
    svc.fit(X, y)
    if save:
        with open('model.pkl', 'wb') as f:
            pickle.dump(svc, f)
    return str(svc.score(X, y))


demo = gr.Interface(fn=svc_demo,
                    title='Iris training',
                    description='调整参数训练模型',
                    inputs=[
                        gr.inputs.Slider(minimum=0, maximum=10, step=0.1, default=1.0),
                        gr.inputs.Dropdown(['linear', 'poly', 'rbf', 'sigmoid'], default='rbf'),
                        gr.inputs.Slider(minimum=1, maximum=10, step=1, default=3),
                        gr.inputs.Dropdown([True, False], default=False)
                    ],
                    outputs="text")


def svc_predict(iris_data):
    with open('model.pkl', 'rb') as f:
        model = pickle.load(f)
    return str(model.predict(iris_data))


predict = gr.Interface(fn=svc_predict,
                       title='Iris predict',
                       description='输入想要预测的数据',
                       inputs=[
                           gr.inputs.Dataframe(type='numpy', row_count=1, col_count=4, datatype='number',
                                               default=[[5.1, 3.5, 1.4, 0.2]])
                       ],
                       outputs='text')

if __name__ == "__main__":
    # demo.launch()
    predict.launch()
